package com.lwc.backendsinglestructure.llm.vector;

import com.lwc.backendsinglestructure.llm.embedding.BasicEmbeddingModel;
import io.milvus.param.R;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import org.springframework.beans.factory.annotation.Autowired;

import java.util.List;

/**
 * 抽象的 Milvus 服务
 */
public abstract class AbstractMilvusService<T> {

    @Autowired
    protected MilvusClientV2 clientV2;
    @Autowired
    protected BasicEmbeddingModel basicModelClient;
    /**
     * 默认查询相似度最高的1条
     */
    private final static int DEFAULT_TOP_K = 1;

    private final static int DEFAULT_SCORE = 0;
    /**
     * 单条文本 embedding（文本转向量）
     *
     * @param text
     * @return
     */
    protected float[] embed(String text) {
        float[] embed = basicModelClient.embed(text);
        // TODO 判空
        return embed;
    }

    /**
     * 批量文本 embedding（文本转向量）
     *
     * @param texts
     * @return
     */
    protected List<float[]> embed(List<String> texts) {
        List<float[]> embed = basicModelClient.embed(texts);
        // TODO 判空
        return embed;
    }


    /**
     * 创建集合
     * @return
     */
    private void createCollection(){
        CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                .collectionName(collcetionName())
                .collectionSchema(schema())
                .indexParams(indexParams())
                .build();

        clientV2.createCollection(customizedSetupReq1);

    }

    /**
     * 表结构，由实现类实现，每个表结可能不一样
     * @return
     */
    protected abstract CreateCollectionReq.CollectionSchema schema ();

    /**
     * 索引结构
     * @return
     */
    protected abstract List<IndexParam> indexParams();


    /**
     * 删除 Collection
     * @return
     */
    private void dropCollection(){
        DropCollectionReq dropQuickSetupParam = DropCollectionReq.builder()
                .collectionName(collcetionName())
                .build();
        clientV2.dropCollection(dropQuickSetupParam);
    };

    /**
     * 集合名称，类似RDBMS中的数据库概念
     *
     * 明明定义由具体实现类来定义
     * @return
     */
    protected abstract String collcetionName();

    /**
     * 基于相似度的 TopK 条数据搜索
     *
     * @param query 用户查询
     * @param topK
     * @return
     */
    public abstract List<R> search(String query, int topK, float gtScore);

    public List<R>  search(String query, int topK) {
        return this.search(query, topK, DEFAULT_SCORE);
    }

    public List<R>  search(String query) {

        return this.search(query, DEFAULT_TOP_K);
    }

    /**
     * 向量纬度
     *
     * @return
     */
    protected abstract int dimension();

    /**
     * 数据写入
     *
     * @param data
     * @return
     */
    public abstract boolean writeCollection(T data);

}
